#ifndef _HASH_TREE_H_
#define  _HASH_TREE_H_

#include "Global.h"
#include "map_utility.h"
#include "Kmer.h"
#include "array.h"
#include "SplayMap.h"
#include <cassert>
using namespace std;

// K is CompressedKmer
// T is the affiliated data type
// M is the inner map type
template<typename T, typename M = SplayMap<array<char, INNER_KEY_LENGTH>, T, PartialKmerLess> >
class HashTree
{
public:
	typedef CompressedKmer key_type;
	typedef T data_type;
	typedef pair<const key_type, data_type> value_type;

	typedef array<char, INNER_KEY_LENGTH> map_key_t;

public:
	typedef M map_t;
protected:
	typedef typename map_t::iterator map_iterator_t;
	typedef typename map_t::const_iterator const_map_iterator_t;
private:
	class SecondaryTree {
	private:
		map_t m_map;
	public:
		//default constructor and destructor are ok.
		pair<map_iterator_t, bool> insert(const value_type& keyValuePair) {
			map_key_t mapKey;
			memcpy(mapKey.memblock, keyValuePair.first.bases + OUTER_KEY_LENGTH, INNER_KEY_LENGTH);
			return m_map.insert(make_pair(mapKey, keyValuePair.second));
		}

		unsigned int erase(const key_type& key) {
			map_key_t mapKey;
			memcpy(mapKey.memblock, key.bases + OUTER_KEY_LENGTH, INNER_KEY_LENGTH);
			return m_map.erase(mapKey);
		}

		map_iterator_t find(const key_type& key) {
			map_key_t mapKey;
			memcpy(mapKey.memblock, key.bases + OUTER_KEY_LENGTH, INNER_KEY_LENGTH);
			return m_map.find(mapKey);
		}

		const_map_iterator_t find(const key_type& key) const {
			map_key_t mapKey;
			memcpy(mapKey.memblock, key.bases + OUTER_KEY_LENGTH, INNER_KEY_LENGTH);
			return m_map.find(mapKey);
		}

		unsigned int count(const key_type& key) const {
			map_key_t mapKey;
			memcpy(mapKey.memblock, key.bases + OUTER_KEY_LENGTH, INNER_KEY_LENGTH);
			return m_map.count(mapKey);
		}

		map_iterator_t begin() {
			return m_map.begin();
		}

		const_map_iterator_t begin() const {
			return m_map.begin();
		}

		map_iterator_t end() {
			return m_map.end();
		}

		const_map_iterator_t end() const {
			return m_map.end();
		}

		unsigned int size() const {
			return m_map.size();
		}

		bool empty() const {
			return m_map.empty();
		}

		void clear() {
			m_map.clear();
		}

		bool write_nopointer_data(FILE* fp) {
			if (!fwrite(&(m_map.size()), sizeof(unsigned int), 1, fp)) {
				return false;
			}
			for (const_map_iterator_t it = m_map.begin(); it != m_map.end(); ++it) {
				if (!fwrite(&(it->first), sizeof(map_key_t), 1, fp) || !fwrite(&(it->second), sizeof(data_type), 1, fp)) {
					return false;
				}
			}
			return true;
		}

		bool read_nopointer_data(FILE* fp) {
			unsigned int entryCnt;
			if (!fread(&entryCnt, sizeof(unsigned int), 1, fp)) {
				return false;
			}
			for (unsigned int i = 0; i < entryCnt; ++i) {
				map_key_t key;
				data_type data;
				if (!fread(&key, sizeof(key_type), 1, fp) || !fread(&data, sizeof(data_type), 1, fp)) {
					return false;
				}
				m_map.insert(make_pair(key, data));
			}
			return true;
		}
	};

	/**
	 * The reference wrapper.
	 */
	template <class _Ty>
	class MutableReference {
		_Ty* _data ;
	public:
		MutableReference(): _data(NULL) {}
		// allow implicit conversion
		MutableReference(_Ty& _dat): _data(&_dat) {}
		MutableReference(_Ty* _dat): _data(_dat) {}

		void setReference(_Ty& _dat) { _data = &_dat; }
		void setPointer(_Ty* _dat) { _data = _dat; }

		_Ty& getReference() { return *_data; }
		_Ty* getPointer() { return _data; }

		operator _Ty() const {
			return *_data;
		}

		_Ty& operator++() {
			return ++(*_data);
		}

		_Ty operator++(int) {
			return (*_data)++;
		}
	};
public:
	class iterator {
	private:
		SecondaryTree* m_hashTable;
		GUINT32 m_outerPos;
		map_iterator_t m_innerIter;
		map_iterator_t m_innerEndIter;
		GUINT32 m_outerEndPos;
		key_type first;
		data_type* second;
		pair<MutableReference<const key_type>, MutableReference<data_type> > m_pair;

		void buildCompleteKeyValue() {
			//build key
			if (m_outerPos != m_outerEndPos || m_innerEndIter != m_innerIter) {
				for (int i = 0; i < OUTER_KEY_LENGTH; ++i) {
					first.bases[i] = ( m_outerPos >> ((OUTER_KEY_LENGTH - i - 1) << 3) ) & 0xff;
				}
				memcpy(first.bases + OUTER_KEY_LENGTH, m_innerIter->first.memblock, INNER_KEY_LENGTH);
				//build data
				second = &(m_innerIter->second);
				m_pair.first.setReference(first);
				m_pair.second.setReference(*second);
			} // if (outerPos != HASHTREE_SIZE - 1 || m_hashtable[outerPos].end() != m_innerEndIter)
		}
	public:
		iterator(): m_hashTable(NULL), second(NULL) {} // only used for declaration
		iterator(SecondaryTree* hashTable, int outerPos, map_iterator_t innerIter, int outerEndPos):
		m_hashTable(hashTable),	m_outerPos(outerPos), m_innerIter(innerIter),
		m_outerEndPos(outerEndPos), m_innerEndIter(m_hashTable[outerEndPos].end()),
		second(NULL)
		{
			buildCompleteKeyValue();
		}

		iterator(const iterator& other) {
			m_hashTable = other.m_hashTable;
			m_outerPos = other.m_outerPos;
			m_innerIter = other.m_innerIter;
			m_innerEndIter = other.m_innerEndIter;
			m_outerEndPos = other.m_outerEndPos;
			first = other.first;
			second = other.second;
			m_pair.first.setReference(first);
			m_pair.second.setReference(*second);
		}

		// Default destructor is fine.

		iterator& operator=(const iterator& other) {
			m_hashTable = other.m_hashTable;
			m_outerPos = other.m_outerPos;
			m_innerIter = other.m_innerIter;
			m_innerEndIter = other.m_innerEndIter;
			m_outerEndPos = other.m_outerEndPos;
			first = other.first;
			second = other.second;
			m_pair.first.setReference(first);
			m_pair.second.setReference(*second);
			return *this;
		}

		// Happy dereferencer
		pair<MutableReference<const key_type>, MutableReference<data_type> >& operator*() {
			return m_pair; 
		}

		pair<MutableReference<const key_type>, MutableReference<data_type> >* operator->() { return &m_pair; }

		// Arithmetic.
		iterator& operator++()   
		{
			assert(m_outerPos != m_outerEndPos || m_innerIter != m_innerEndIter);
			if (m_hashTable[m_outerPos].end() != (++m_innerIter)) {
				buildCompleteKeyValue();
				return *this;
			} // if (m_hashTable[m_outerPos].end() != (++m_innerIter))
			else if (m_outerPos == m_outerEndPos && m_innerIter == m_innerEndIter) {
				return *this; // Oh! End meet!
			}
			for (++m_outerPos; m_outerPos < m_outerEndPos; ++m_outerPos) {
				if (!m_hashTable[m_outerPos].empty()) { break; }
			}
			m_innerIter = m_hashTable[m_outerPos].begin();
			buildCompleteKeyValue();
			return *this;
		}
		iterator operator++(int) { iterator tmp(*this); ++*this; return tmp; }

		// Comparison.
		bool operator==(const iterator& it) const {
			return m_hashTable == it.m_hashTable &&
			m_outerPos == it.m_outerPos && m_innerIter == it.m_innerIter &&
			m_innerEndIter == it.m_innerEndIter && m_outerEndPos == m_outerEndPos;
		}
		bool operator!=(const iterator& it) const { return !operator==(it); }
	};

	class const_iterator {
	private:
		const SecondaryTree* m_hashTable;
		GUINT32 m_outerPos;
		GUINT32 m_outerEndPos;
		const_map_iterator_t m_innerIter;
		const_map_iterator_t m_innerEndIter;
		key_type first;
		data_type* second;
		pair<MutableReference<const key_type>, MutableReference<data_type> > m_pair;

		void buildCompleteKeyValue() {
			//build key
			if (m_outerPos != m_outerEndPos || m_innerEndIter != m_innerIter) {
				for (int i = 0; i < OUTER_KEY_LENGTH; ++i) {
					first.bases[i] = ( m_outerPos >> ((OUTER_KEY_LENGTH - i - 1) << 3) ) & 0xff;
				}
				memcpy(first.bases + OUTER_KEY_LENGTH, m_innerIter->first.memblock, INNER_KEY_LENGTH);
				//build data
				second = &(m_innerIter->second);
				m_pair.first.setReference(first);
				m_pair.second.setReference(*second);
			} // if (outerPos != HASHTREE_SIZE - 1 || m_hashtable[outerPos].end() != m_innerEndIter)
		}
	public:
		const_iterator(): m_hashTable(NULL), second(NULL) {} // only used for declaration
		const_iterator(const SecondaryTree* hashTable, int outerPos, const_map_iterator_t innerIter, int outerEndPos):
			m_hashTable(hashTable),	m_outerPos(outerPos), m_innerIter(innerIter),
			m_outerEndPos(outerEndPos), m_innerEndIter(m_hashTable[outerEndPos].end()),
			second(NULL)
		{
			buildCompleteKeyValue();
		}

		const_iterator(const const_iterator& other) {
			m_hashTable = other.m_hashTable;
			m_outerPos = other.m_outerPos;
			m_innerIter = other.m_innerIter;
			m_innerEndIter = other.m_innerEndIter;
			m_outerEndPos = other.m_outerEndPos;
			first = other.first;
			second = other.second;
			m_pair.first.setReference(first);
			m_pair.second.setReference(*second);
		}

		const_iterator(const iterator& other) {
			m_hashTable = other.m_hashTable;
			m_outerPos = other.m_outerPos;
			m_innerIter = other.m_innerIter;
			m_innerEndIter = other.m_innerEndIter;
			m_outerEndPos = other.m_outerEndPos;
			first = other.first;
			second = other.second;
			m_pair.first.setReference(first);
			m_pair.second.setReference(*second);
		}

		// Default destructor is fine.

		const_iterator& operator=(const const_iterator& other) {
			m_hashTable = other.m_hashTable;
			m_outerPos = other.m_outerPos;
			m_innerIter = other.m_innerIter;
			m_innerEndIter = other.m_innerEndIter;
			m_outerEndPos = other.m_outerEndPos;
			first = other.first;
			second = other.second;
			m_pair.first.setReference(first);
			m_pair.second.setReference(*second);
			return *this;
		}

		const_iterator& operator=(const iterator& other) {
			m_hashTable = other.m_hashTable;
			m_outerPos = other.m_outerPos;
			m_innerIter = other.m_innerIter;
			m_innerEndIter = other.m_innerEndIter;
			m_outerEndPos = other.m_outerEndPos;
			first = other.first;
			second = other.second;
			m_pair.first.setReference(first);
			m_pair.second.setReference(*second);
			return *this;
		}

		// Happy dereferencer
		const pair<const key_type&, data_type&>& operator*() const
		{
			return m_pair; 
		}

		const pair<const key_type&, data_type&>* operator->() const { return &m_pair; }

		// Arithmetic.
		const_iterator& operator++()   
		{
			assert(m_outerPos != m_outerEndPos || m_innerIter != m_innerEndIter);
			if (m_hashTable[m_outerPos].end() != (++m_innerIter)) {
				buildCompleteKeyValue();
				return *this;
			} // if (m_hashTable[m_outerPos].end() != (++m_innerIter))
			else if (m_outerPos == m_outerEndPos && m_innerIter == m_innerEndIter) {
				return *this; // Oh! End meet!
			}
			for (++m_outerPos; m_outerPos < m_outerEndPos; ++m_outerPos) {
				if (!m_hashTable[m_outerPos].empty()) { break; }
			}
			m_innerIter = m_hashTable[m_outerPos].begin();
			buildCompleteKeyValue();
			return *this;
		}
		const_iterator operator++(int) { const_iterator tmp(*this); ++*this; return tmp; }

		// Comparison.
		bool operator==(const const_iterator& it) const {
			return m_hashTable == it.m_hashTable &&
				m_outerPos == it.m_outerPos && m_innerIter == it.m_innerIter &&
				m_innerEndIter == it.m_innerEndIter && m_outerEndPos == m_outerEndPos;
		}
		bool operator!=(const const_iterator& it) const { return !operator==(); }
	};
	
private:
	SecondaryTree m_hashtable[HASHTREE_SIZE];
	int m_beginPos, m_endPos;
	unsigned int m_size;

	unsigned int getHashKey(const key_type& key) {
		unsigned int hashKey = 0;
		for (int i = 0; i < OUTER_KEY_LENGTH; ++i) {
			hashKey |= key.bases[i] << ( (OUTER_KEY_LENGTH - i - 1) << 3);
		}
		return hashKey;
	}
public:
	iterator begin() {
		return iterator(m_hashtable, m_beginPos,  m_hashtable[m_beginPos].begin(), m_endPos);
	}

	const_iterator begin() const {
		return const_iterator(m_hashtable, m_beginPos,  m_hashtable[m_beginPos].begin(), m_endPos);
	}

	iterator end() {
		return iterator(m_hashtable, m_endPos,  m_hashtable[m_endPos].end(), m_endPos);
	}

	const_iterator end() const {
		return const_iterator(m_hashtable, m_endPos,  m_hashtable[m_endPos].end(), m_endPos);
	}
	
	HashTree(): m_beginPos(0), m_endPos(0), m_size(0) {}
	~HashTree() {}
	
	void clear() {
		for (int i = 0; i < HASHTREE_SIZE; ++i) {
			m_hashtable[i].clear();
		}
		m_size = 0;
	}

	unsigned int size() { return m_size; }

	template<class _Cd>
	size_t bucketCount(size_t allBucketCount[HASHTREE_SIZE], _Cd& conditionor) {
		size_t total = 0;
		memset(allBucketCount, 0, sizeof(size_t) * HASHTREE_SIZE);
		for (int i = 0; i < HASHTREE_SIZE; ++i) {
			for (typename map_t::iterator it = m_hashtable[i].begin(); it != m_hashtable[i].end(); ++it) {
				if (conditionor(it->second)) {
					++allBucketCount[i];
					++total;
				}
			}
		} // for (int i = 0; i < HASHTREE_SIZE; ++i)
		return total;
	}

	size_t bucketCount(size_t allBucketCount[HASHTREE_SIZE]) {
		for (int i = 0; i < HASHTREE_SIZE; ++i) {
			allBucketCount[i] = m_hashtable[i].size();
		}
		return m_size;
	}

	bool empty() {
		return !size();
	}

	iterator find(const key_type& key) {
		unsigned int hashKey = getHashKey(key);
		map_iterator_t res = m_hashtable[hashKey].find(key);
		if (res != m_hashtable[hashKey].end()) {
			return iterator(m_hashtable, hashKey, res, m_endPos);
		}
		return end();
	}

	const_iterator find(const key_type& key) const {
		unsigned int hashKey = getHashKey(key);
		const_map_iterator_t res = m_hashtable[hashKey].find(key);
		if (res != m_hashtable[hashKey].end()) {
			return iterator(m_hashtable, hashKey, res, m_endPos);
		}
		return end();
	}

	data_type& operator[](const key_type& key) {
		unsigned int hashKey = getHashKey(key);
		const_map_iterator_t res = m_hashtable[hashKey].find(key);
		if (res != m_hashtable[hashKey].end()) {
			return res->second;
		}
		// insert default value
		pair<map_iterator_t, bool> ins = m_hashtable[hashKey].insert(make_pair(key, data_type()));
		if (ins.second) {
			++m_size;
			// exception is for initial begin value 0
			if (hashKey < m_beginPos || (m_beginPos == 0 && m_hashtable[m_beginPos].empty())) {
				m_beginPos = hashKey;
			}
			if (hashKey > m_endPos) {
				m_endPos = hashKey;
			}
		}
		return ins.first->second;
	}

	unsigned int count(const key_type& key) {
		if (find(key) != end()) {
			return 1;
		}
		return 0;
	}

	pair<iterator, bool> insert(const value_type& keyDataPair) {
		unsigned int hashKey = getHashKey(keyDataPair.first);
		pair<map_iterator_t, bool> ins = m_hashtable[hashKey].insert(keyDataPair);
		if (ins.second) {
			++m_size;
			if (hashKey < m_beginPos || (m_beginPos == 0 && m_hashtable[m_beginPos].empty())) {
				m_beginPos = hashKey;
			}
			if (hashKey > m_endPos) {
				m_endPos = hashKey;
			}
		}
		return make_pair(iterator(m_hashtable, hashKey, ins.first, m_endPos), ins.second);
	}

	unsigned int erase(const key_type& key) {
		unsigned int hashKey = getHashKey(key);
		unsigned int erased =  m_hashtable[hashKey].erase(key);
		if (erased) {
			--m_size;
			if (hashKey == m_beginPos) {
				for (; m_beginPos < HASHTREE_SIZE; ++m_beginPos) {
					if (!m_hashtable[m_beginPos].empty()) { break; }
					if (m_beginPos == HASHTREE_SIZE) { m_beginPos = 0; }
				} // if (hashKey == m_beginPos) {for (; m_beginPos < HASHTREE_SIZE; ++m_beginPos)
			} // if (hashKey == m_beginPos)
			if (hashKey == m_endPos) {
				for (; m_endPos > 0; --m_endPos) {
					if (!m_hashtable[m_endPos].empty()) { break; }
				}
			} // if (hashKey == m_endPos)
		}
		return erased;
	}

	bool write_data(FILE *fp) {
		if (!fwrite(&m_beginPos, sizeof(int), 1, fp) || !fwrite(&m_endPos, sizeof(int), 1, fp)
			|| !fwrite(&m_size, sizeof(unsigned int), 1, fp)) {
				return false;
		}
		for (int i = 0; i < HASHTREE_SIZE; ++i) {
			if (!m_hashtable[i].write_data(fp)) {
				return false;
			}
		} // for (int i = 0; i < HASHTREE_SIZE; ++i)
		return true;
	};

	bool read_data(FILE *fp) {
		if (!fread(&m_beginPos, sizeof(int), 1, fp) || !fread(&m_endPos, sizeof(int), 1, fp)
			|| !fread(&m_size, sizeof(unsigned int), 1, fp)) {
				return false;
		}
		for (int i = 0; i < HASHTREE_SIZE; ++i) {
			if (!m_hashtable[i].read_data(fp)) {
				return false;
			}
		} // for (int i = 0; i < HASHTREE_SIZE; ++i)
		return true;
	}
};

#endif